package com.abyss;

import org.apache.flink.api.common.functions.RichMapFunction;
import org.apache.flink.api.java.ExecutionEnvironment;
import org.apache.flink.api.java.operators.DataSource;
import org.apache.flink.api.java.operators.MapOperator;
import org.apache.flink.api.java.tuple.Tuple2;
import org.apache.flink.api.java.tuple.Tuple3;
import org.apache.flink.configuration.Configuration;

import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
 * 广播变量
 */
public class BroadcastDemo {

    public static void main(String[] args) throws Exception {
        // 1. Env
        ExecutionEnvironment env = ExecutionEnvironment.getExecutionEnvironment();

        // 2. source
        // 2.1 学生信息source
        DataSource<Tuple2<Integer, String>> studentSource = env.fromElements(
                Tuple2.of(1, "王大锤"),
                Tuple2.of(2, "潇潇"),
                Tuple2.of(3, "甜甜")
        );
        // 2.2 学生成绩信息source
        DataSource<Tuple3<Integer, String, Integer>> scoreSource = env.fromElements(
                Tuple3.of(1, "数据结构", 99),
                Tuple3.of(2, "英语", 100),
                Tuple3.of(3, "C++", 96),
                Tuple3.of(5, "Java", 97),
                Tuple3.of(3, "Scala", 100)
        );

        /*
        广播变量：
        1. 先注册： withBroadcastSet方法来注册
        2. 使用它: runtimeContext对象来获得广播变量的对象
         */
        MapOperator<Tuple3<Integer, String, Integer>, Tuple3<String, String, Integer>> result = scoreSource.map(new RichMapFunction<Tuple3<Integer, String, Integer>, Tuple3<String, String, Integer>>() {
            private Map<Integer, String> stuMap = new HashMap<>();

            @Override
            public void open(Configuration parameters) throws Exception {
                super.open(parameters);
                List<Tuple2<Integer, String>> studentList = getRuntimeContext().getBroadcastVariable("student");
                for (Tuple2<Integer, String> stu : studentList) {
                    stuMap.put(stu.f0, stu.f1);
                }
            }

            @Override
            public void close() throws Exception {
                super.close();
            }

            @Override
            public Tuple3<String, String, Integer> map(Tuple3<Integer, String, Integer> value) throws Exception {
                int stuId = value.f0;
                String stuName = this.stuMap.getOrDefault(stuId, "未知的学生名称");
                return Tuple3.of(stuName, value.f1, value.f2);
            }
        }).withBroadcastSet(studentSource, "student");
        result.print();
    }
}
